//! Code generation from the Python low-level IR.

use stencil_codegen::{CodeBuffer, Error};
use stencil_tysys::{Endianness, IntEncoding, IntType, MonOp};

use super::ir::{
    Exception, FnSignature, Lvalue, LvalueRef, PyBlock, PyClassDef, PyDef, PyExpr, PyFnDef,
    PyImport, PyModule, PyStatement,
};

pub fn emit_module(m: &PyModule, buf: &mut CodeBuffer) -> Result<(), Error> {
    if let Some(comment) = m.initial_comment.as_ref() {
        emit_comment(&comment, buf)?;
        buf.emit_space();
    }

    for def in &m.defs {
        emit_def(def, buf)?;
    }

    Ok(())
}

fn emit_comment(comment: &str, buf: &mut CodeBuffer) -> Result<(), Error> {
    for line in comment.lines() {
        buf.emit_fmt(format_args!("# {line}"));
    }
    Ok(())
}

fn emit_doc(doc: &str, buf: &mut CodeBuffer) -> Result<(), Error> {
    buf.emit_line("\"\"\"");
    for line in doc.lines() {
        buf.emit_line(line);
    }
    buf.emit_line("\"\"\"");
    Ok(())
}

fn emit_def(d: &PyDef, buf: &mut CodeBuffer) -> Result<(), Error> {
    match d {
        PyDef::Import(imp) => emit_import(imp, buf),
        PyDef::Class(cls) => emit_classdef(cls, buf),
        _ => unimplemented!(),
    }
}

fn emit_import(imp: &PyImport, buf: &mut CodeBuffer) -> Result<(), Error> {
    match imp {
        PyImport::Simple(mspec) => {
            let imp_path = mspec.full();
            buf.emit_fmt(format_args!("import {imp_path}"));
            Ok(())
        }
        PyImport::Specific(mspec, items) => {
            let imp_path = mspec.full();
            let items = concat_intersperse_strs(items, ", ");
            buf.emit_fmt(format_args!("from {imp_path} import {items}"));
            Ok(())
        }
    }
}

fn emit_classdef(cls: &PyClassDef, buf: &mut CodeBuffer) -> Result<(), Error> {
    if cls.is_dataclass {
        buf.emit_line("@dataclass");
    }

    let name = cls.name.as_str();
    buf.emit_fmt(format_args!("class {name}():"));

    buf.push_indent();

    if let Some(doc) = &cls.doc {
        emit_doc(doc, buf)?;
    }

    for f in &cls.fields {
        let name = f.name.as_str();
        let ty_str = f.ty.to_python();
        buf.emit_fmt(format_args!("{name}: {ty_str}"));
    }

    if !cls.fields.is_empty() {
        buf.emit_space();
    }

    for f in &cls.functions {
        emit_function(f, buf)?;
    }

    buf.pop_indent();

    Ok(())
}

fn emit_function(fdef: &PyFnDef, buf: &mut CodeBuffer) -> Result<(), Error> {
    if fdef.is_static {
        buf.emit_line("@staticmethod");
    }

    let name = fdef.name.as_str();
    let sig_str = fmt_signature(&fdef.sig);
    buf.emit_fmt(format_args!("def {name}{sig_str}:"));

    buf.push_indent();

    if let Some(doc) = &fdef.doc {
        emit_doc(doc, buf)?;
        buf.emit_space();
    }

    emit_block(&fdef.body, fdef, buf)?;

    buf.pop_indent();
    buf.emit_space();

    Ok(())
}

fn emit_block(block: &PyBlock, fdef: &PyFnDef, buf: &mut CodeBuffer) -> Result<(), Error> {
    for stmt in &block.statements {
        emit_statement(stmt, fdef, buf)?;
    }
    Ok(())
}

fn emit_statement(stmt: &PyStatement, fdef: &PyFnDef, buf: &mut CodeBuffer) -> Result<(), Error> {
    match stmt {
        PyStatement::Comment(m) => {
            emit_comment(m, buf)?;
        }

        PyStatement::Assign(target, expr) => {
            let lvc = gen_lvalue(target)?;
            let exprc = gen_expr(expr)?;
            buf.emit_fmt(format_args!("{lvc} = {exprc}"));
        }

        PyStatement::AssignIndex(target, index, value) => {
            let lvc = gen_lvalue(target)?;
            let indexc = gen_expr(index)?;
            let valuec = gen_expr(value)?;
            buf.emit_fmt(format_args!("{lvc}[{indexc}] = {valuec}"));
        }

        PyStatement::AssignBufInt(target, off, value, ienc) => {
            let lv_str = gen_lvalue(target)?;
            let off_str = gen_expr(off)?;
            let val_str = gen_expr(value)?;

            if ienc.ty() == IntType::U8 {
                buf.emit_fmt(format_args!("{lv_str}[{off_str}] = {val_str}"));
            } else {
                let len = ienc.ty().byte_len();
                let endian_str = match ienc.endianness() {
                    Endianness::Big => "'big'",
                    Endianness::Little => "'little'",
                };

                buf.emit_fmt(format_args!("{lv_str}[{off_str}:{off_str}+{len}] = int.to_bytes({val_str}, length={len}, byteorder={endian_str})"));
            }
        }

        PyStatement::Return(expr) => {
            let exprc = gen_expr(expr)?;
            buf.emit_fmt(format_args!("return {exprc}"));
        }

        PyStatement::ReturnNone => {
            buf.emit_line("return");
        }

        PyStatement::Raise(exc) => {
            emit_raise(exc, buf)?;
        }

        PyStatement::If(expr, body) => {
            let exprc = gen_expr(expr)?;
            buf.emit_fmt(format_args!("if {exprc}:"));
            buf.push_indent();
            emit_block(&body, fdef, buf)?;
            buf.pop_indent();
        }

        PyStatement::IfElse(expr, then_body, else_body) => {
            let exprc = gen_expr(expr)?;
            buf.emit_fmt(format_args!("if {exprc}:"));
            buf.push_indent();
            emit_block(&then_body, fdef, buf)?;
            buf.pop_indent();
            buf.emit_line("else:");
            buf.push_indent();
            emit_block(&else_body, fdef, buf)?;
            buf.pop_indent();
        }

        PyStatement::While(expr, body) => {
            let exprc = gen_expr(expr)?;
            buf.emit_fmt(format_args!("while {exprc}:"));
            buf.push_indent();
            emit_block(&body, fdef, buf)?;
            buf.pop_indent();
        }

        PyStatement::ForRange(ident, end_expr, body) => {
            let exprc = gen_expr(end_expr)?;
            buf.emit_fmt(format_args!("for {} in range({exprc}):", ident.as_str()));
            buf.push_indent();
            emit_block(&body, fdef, buf)?;
            buf.pop_indent();
        }
    }

    Ok(())
}

fn emit_raise(exc: &Exception, buf: &mut CodeBuffer) -> Result<(), Error> {
    match exc {
        Exception::ValueError(m) => {
            buf.emit_fmt(format_args!("raise ValueError({m:?})"));
        }
    }

    Ok(())
}

fn gen_expr(expr: &PyExpr) -> Result<String, Error> {
    Ok(match expr {
        PyExpr::ClassSelf => "self".to_string(),
        PyExpr::Literal(lit) => lit.to_python(),
        PyExpr::Ident(ident) => ident.as_str().to_owned(),
        PyExpr::Ternary(ce, te, ee) => {
            let cc = gen_expr(ce)?;
            let tc = gen_expr(te)?;
            let ec = gen_expr(ee)?;
            format!("{tc} if {cc} else {ec}")
        }
        PyExpr::Cast(ty, expr) => format!("{}({})", ty.to_python(), gen_expr(expr)?),
        PyExpr::Monop(op, expr) => {
            let v = match op {
                MonOp::Neg => "!",
                MonOp::BNot => "~",
                MonOp::LNot => "not ",
            };
            format!("{v}{}", gen_expr(expr)?)
        }
        PyExpr::Binop(op, lhs, rhs) => {
            // FIXME does not handle LOr and LAnd properly since Python uses "or" and "and"
            let v = op.to_operator_str();
            format!("({} {v} {})", gen_expr(lhs)?, gen_expr(rhs)?)
        }
        PyExpr::Dot(expr, member) => format!("{}.{}", gen_expr(expr)?, member.as_str()),
        PyExpr::Index(expr, idx) => format!("{}[{}]", gen_expr(expr)?, gen_expr(idx)?),
        PyExpr::Len(expr) => format!("len({})", gen_expr(expr)?),
        PyExpr::BytebufReadSlice(target, start, end) => gen_slice_read(target, start, end)?,
        PyExpr::BytebufReadInt(target, off, encoding) => gen_int_read(target, off, *encoding)?,
    })
}

fn gen_slice_read(_target: &PyExpr, _start: &PyExpr, _end: &PyExpr) -> Result<String, Error> {
    // We don't generate these yet.
    unimplemented!()
}

fn gen_int_read(target: &PyExpr, off: &PyExpr, ienc: IntEncoding) -> Result<String, Error> {
    let tc = gen_expr(target)?;
    let oc = gen_expr(off)?;

    let ity = ienc.ty();
    let len = ity.byte_len();
    let signed_str = if ity.is_signed() { "True" } else { "False" };

    let byteord_str = format!("\"{}\"", ienc.endianness().lower_str());

    Ok(if ity == IntType::U8 {
        format!("{tc}[{oc}]")
    } else {
        format!(
            "int.from_bytes({tc}[({oc}):({oc})+{len}], byteorder={byteord_str}, signed={signed_str})"
        )
    })
}

fn gen_lvalue(lvalue: &Lvalue) -> Result<String, Error> {
    Ok(match lvalue {
        Lvalue::Ref(lvr) => gen_lvalue_ref(lvr)?,
        Lvalue::RefField(lvr, member) => format!("{}.{}", gen_lvalue_ref(lvr)?, member.as_str()),
        /*Lvalue::RefIndex(lvr, expr) => format!("{}[{}]", gen_lvalue_ref(lvr)?, gen_expr(expr)?),
        Lvalue::RefFieldIndex(lvr, member, expr) => format!(
            "{}.{}[{}]",
            gen_lvalue_ref(lvr)?,
            member.as_str(),
            gen_expr(expr)?
        ),*/
    })
}

fn gen_lvalue_ref(lvr: &LvalueRef) -> Result<String, Error> {
    Ok(match lvr {
        LvalueRef::ClassSelf => "self".to_owned(),
        LvalueRef::Local(l) => format!("{}", l.as_str()),
    })
}

/// Generates a function signature.
fn fmt_signature(sig: &FnSignature) -> String {
    let mut arg_buf = String::new();

    if sig.is_self {
        arg_buf.push_str("self");
    }

    for arg in &sig.args {
        if !arg_buf.is_empty() {
            arg_buf.push_str(", ");
        }

        let name = arg.name.as_str();
        let arg = match (&arg.ty, &arg.val) {
            (Some(ty), Some(val)) => format!("{name}: {} = {}", ty.to_python(), val.to_python()),
            (Some(ty), None) => format!("{name}: {}", ty.to_python()),
            (None, Some(val)) => format!("{name} = {}", val.to_python()),
            _ => {
                arg_buf.push_str(name);
                continue;
            }
        };

        arg_buf.push_str(&arg);
    }

    if let Some(rty) = &sig.ret_ty {
        format!("({arg_buf}) -> {}", rty.to_python())
    } else {
        format!("({arg_buf})")
    }
}

fn concat_intersperse_strs<S: AsRef<str>>(
    parts: impl IntoIterator<Item = S>,
    inter: &str,
) -> String {
    let mut s = String::new();
    let mut parts_iter = parts.into_iter();
    while let Some(item) = parts_iter.next() {
        if !s.is_empty() {
            s.push_str(inter);
        }

        s.push_str(item.as_ref());
    }
    s
}
